"""
Run a rest API exposing the yolov5s object detection model
"""
import argparse
import io
import numpy as np
import cv2
import json

import base64

import torch
from flask import Flask, request, make_response
from PIL import Image
from flask_cors import *

app = Flask(__name__)

DETECTION_URL = "/v1/object-detection/yolov5s"


@app.route(DETECTION_URL, methods=["POST"])
@cross_origin(supports_credentials=True)
def predict():
    if not request.method == "POST":
        return

    # 处理图片
    # if request.data:
    #     img = cv2.imdecode(np.frombuffer(request.data, dtype=np.uint8), cv2.IMREAD_COLOR)
    #     results = model(img)  # reduce size=320 for faster inference
    #     return cv2.imencode(".jpg", results)[1].tobytes()
    # return results.pandas().xyxy[0].to_json(orient="records")

    if request.files.get("image"):
        image_file = request.files["image"]
        image_bytes = image_file.read()
        img = Image.open(io.BytesIO(image_bytes))

        results = model(img, size=640)  # reduce size=320 for faster inference
        info = results.__str__()
        # info = json.dumps(info)
        # info = json.loads(results.pandas().xyxy[0].to_json(orient="records"))
        results = results.render()[0]
        # 颜色转化
        results = cv2.cvtColor(results, cv2.COLOR_BGR2RGB)

        img_raw = cv2.imencode(".jpg", results)[1].tobytes()

        # img_raw = cv2.imdecode(np.frombuffer(img_raw, dtype=np.uint8), cv2.IMREAD_COLOR)
        base64_img_str = base64.b64encode(img_raw).decode()
        # base64_img_str = base64.b64encode(img_raw)

        # print(base64_img)
        data = {
            "info": info,
            "img_base64": base64_img_str
        }

        # resp = make_response("Hello")
        # return resp

        return json.dumps(data), 200

        # return cv2.imencode(".jpg", results)[1].tobytes()

        # return results.pandas().xyxy[0].to_json(orient="records")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Flask API exposing YOLOv5 model")
    parser.add_argument("--port", default=5000, type=int, help="port number")
    args = parser.parse_args()

    # 参数（模型位置，custom写死表本地，模型名字，local写死）
    model = torch.hub.load('./', 'custom', path="runs/train/exp11/weights/best.pt",
                           source="local")  # force_reload to recache
    app.run(host="0.0.0.0", port=args.port)  # debug=True causes Restarting with stat
